import numpy as np
import tensorflow as tf


def h_sigmoid(x):
    output = tf.keras.layers.Activation('hard_sigmoid')(x)

    return output


def h_swish(x):
    output = x * h_sigmoid(x)

    return output


def Squeeze_excitation_layer(x):
    inputs = x
    squeeze = inputs.shape[-1] / 2
    excitation = inputs.shape[-1]
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(squeeze)(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.Dense(excitation)(x)
    x = h_sigmoid(x)
    x = tf.keras.layers.Reshape((1, 1, excitation))(x)
    x = inputs * x

    return x


def BottleNeck(inputs, exp_size, out_size, kernel_size, strides, is_se_existing, activation):
    x = tf.keras.layers.Conv2D(filters=exp_size,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    if activation == "HS":
        x = h_swish(x)
    elif activation == "RE":
        x = tf.keras.layers.Activation(tf.nn.relu6)(x)
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=kernel_size,
                                        strides=strides,
                                        padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    if activation == "HS":
        x = h_swish(x)
    elif activation == "RE":
        x = tf.keras.layers.Activation(tf.nn.relu6)(x)
    if is_se_existing:
        x = Squeeze_excitation_layer(x)
    x = tf.keras.layers.Conv2D(filters=out_size,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation(tf.keras.activations.linear)(x)
    if strides == 1 and inputs.shape[-1] == out_size:
        print("x:", x)
        print("inputs:", inputs)
        x = tf.keras.layers.add([x, inputs])

    return x


def MobileNetV3Large(inputs, classes=10):
    x = tf.keras.layers.Conv2D(filters=16,
                               kernel_size=(3, 3),
                               strides=2,
                               padding="same")(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    x = h_swish(x)

    x = BottleNeck(x, exp_size=16, out_size=16, kernel_size=3, strides=1, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=64, out_size=24, kernel_size=3, strides=2, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=72, out_size=24, kernel_size=3, strides=1, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=72, out_size=40, kernel_size=5, strides=2, is_se_existing=True, activation="RE")
    x = BottleNeck(x, exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="RE")
    x = BottleNeck(x, exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="RE")
    x = BottleNeck(x, exp_size=240, out_size=80, kernel_size=3, strides=2, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=200, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=480, out_size=112, kernel_size=3, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=672, out_size=112, kernel_size=3, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=672, out_size=160, kernel_size=5, strides=2, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True, activation="HS")

    x = tf.keras.layers.Conv2D(filters=960,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = h_swish(x)
    x = tf.keras.layers.AveragePooling2D(pool_size=(7, 7), strides=1)(x)
    x = tf.keras.layers.Conv2D(filters=1280,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = h_swish(x)
    x = tf.keras.layers.Conv2D(filters=classes,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same",
                               activation=tf.keras.activations.softmax)(x)

    return x


def MobileNetV3Small(inputs, classes=10):
    x = tf.keras.layers.Conv2D(filters=16,
                               kernel_size=(3, 3),
                               strides=2,
                               padding="same")(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    x = h_swish(x)

    x = BottleNeck(x, exp_size=16, out_size=16, kernel_size=3, strides=2, is_se_existing=True, activation="RE")
    x = BottleNeck(x, exp_size=72, out_size=24, kernel_size=3, strides=2, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=88, out_size=24, kernel_size=3, strides=1, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=96, out_size=40, kernel_size=5, strides=2, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=240, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=240, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=120, out_size=48, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=144, out_size=48, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=288, out_size=96, kernel_size=5, strides=2, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=576, out_size=96, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=576, out_size=96, kernel_size=5, strides=1, is_se_existing=True, activation="HS")

    x = tf.keras.layers.Conv2D(filters=576,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = h_swish(x)

    x = tf.keras.layers.AveragePooling2D(pool_size=(7, 7), strides=1)(x)
    x = tf.keras.layers.Conv2D(filters=1280,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = h_swish(x)
    x = tf.keras.layers.Conv2D(filters=classes,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same",
                               activation=tf.keras.activations.softmax)(x)

    return x


# inputs = np.zeros((1, 224, 224, 3), np.float32)
# print(inputs.shape)
# print(inputs)
# output_large = MobileNetV3Large(inputs)
# output_small = MobileNetV3Small(inputs)

cifar10 = tf.keras.datasets.cifar10
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
# print("x_train.shape", x_train.shape)
# 给数据增加一个维度，使数据和网络结构匹配
# x_train = tf.broadcast_to(x_train, (5000, 32, 32, 3))
# x_test = tf.broadcast_to(x_test, (5000, 32, 32, 3))
x_train = tf.keras.backend.resize_images(x_train[:500], 7, 7, "channels_last", "nearest")
# x_train = x_train.numpy()
# print(type(x_train))
output_large = MobileNetV3Large(x_train)
#
# print(output_large)
print("output_large:", output_large)
